#/*************************************************************************
#    > File Name: Lib/Parallel_Operation.py
#    > Author: Yan Wang
#    > Mail: wangyan@imnu.edu.cn
#    > Created Time: 2022年05月26日 星期四 15时23分41秒
# ************************************************************************/
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from Library import *
from File_Operation import *

def PBS_Running(param,pbs_template="../Examples/test_input/pbs_template.pbs"):
    print("in PBS_Running, the input param is ", param )
    for folder in param["folder"]:
        pbs_file=folder+"/run_pbs.pbs"
        Copy_File(pbs_template,pbs_file)
        print("pbs_file",pbs_file)
        linenos=[2,7,8,18]
        strings=["#PBS -N {}".format(folder), 
                "#PBS -e {}/job.error".format(folder),
                "#PBS -o {}/job.log".format(folder),
                "cd {}".format(folder)]
        File_Replace(pbs_file,linenos,strings)
        os.system(param["command"].format(folder))
    return()


def Slurm_Running(param,slurm_template="../Examples/test_input/slurm_template.pbs"):
    print("in Slurm_Running, the input param is ", param )
    for folder in param["folder"]:
        slurm_file=folder+"/run_slurm.pbs"
        Copy_File(slurm_template,slurm_file)
        print("slurm_file",slurm_file)
        linenos=[2,10]
        strings=["#SBATCH -J {}".format(folder), "cd {}".format(folder)]
        File_Replace(slurm_file,linenos,strings)
        os.system(param["command"].format(folder))
    return()


# 进入多个核，并行计算
# 第一个参数为需要运行的函数， 第二个paras为这个函数的参数组成的列表，  paras中每一个元素会开辟一个新核，运行run_func(paras_element)
def Parallel_Running(run_func, paras,record_file="None"):
    start_t = datetime.datetime.now()   #.datetime.now()获取指定日期格式的系统的本地时间
    num_cores = int(mp.cpu_count())
    # 使用的核数，不必超过需要的参数个数
    if num_cores> len(paras):
        num_cores=len(paras)
    print("local PC has " + str(num_cores) + " cores.")
    pool = mp.Pool(num_cores)

    print("The input Function is: ", run_func.__name__)
    print("The input params are: \n", paras)

    results=[pool.apply_async(run_func, args = (this_para,)) for this_para in paras]   #pool.apply_async()结果的顺序不能保证与调用的顺序相同
    print("results type ***************",type(results))
    if results is None:
        print("Finish Parallel Running once, and no return value.")
    else:
        #results=[print(p) for p in results]
        results=[print(p.get()) for p in results]


####if record_file != "None":
####	for index_c, content in enumerate(results):
####	    this_file=open(record_file.format(index_c),"w")
####	    this_file.write(content)
####	    this_file.close()

    end_t=datetime.datetime.now()
    elapsed_sec=(end_t - start_t).total_seconds()
    print("Multi processes calculation,  total costs are: " +
            "{:.2f}".format(elapsed_sec) + " second.")
    return()

# 通过SSH执行命令
# 这些节点最好是添加过 host 的授信节点， 不需要再次填写密码
# 对于已授信节点 ssh_para 应存在 "node_name" 的key， 通过name连接ssh
# 如果需要填写密码，则 ssh_para中 不可存在 node_name的key， 通过ip， user_name, passwd 连接ssh
# 命令应存放在 ssh_para 的 "command" key 中
def Run_SSH(ssh_para):
    import paramiko
    ssh = paramiko.SSHClient()   #远程连接服务器，建立一个sshclient对象
    # 允许链接不在 know_hosts文件中的主机
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    # 建立连接
    if "node_name" in ssh_para.keys():
        ssh.connect(ssh_para["node_name"], port=22)
    else:
        ssh.connect(ssh_para["ip"], username=ssh_para["user"], port=22, password=ssh_para["passwd"])
    # 使用这个连接执行命令
    sh_stdin, ssh_stdout, ssh_stderr = ssh.exec_command(ssh_para["command"])
    # 获取输出
    print(ssh_stdout.read())
    # 关闭连接
    ssh.close()

# 使用并行计算的方式，每个核打开一个SSH 连接到集群的一个节点，  每个节点执行一个命令
# 这里command 第一个参数为需要执行的命令， 第二个参数为执行命令使用的节点名称
def Run_Command_with_SSH(ssh_para):

    if "node_name" in ssh_para.keys():
        print(" We will execute the commmand:",ssh_para["command"], "\n in node:",ssh_para["node_name"])
    else:
        print(" We will execute the commmand:",ssh_para["command"], "\n in node:",ssh_para["ip"])

    result = subprocess.run(Run_SSH(ssh_para), capture_output=True, shell=True)  # capture_output=True获stdout和stderr，调用时内部的Popen对象将自动使用stdout=PIPE和stderr = PIPE创建标准输出和标准错误对象；
    return(result)

